from main import setup_parser, merge_config, load_json
from utils.data_manager import DataManager
from torch.utils.data import DataLoader
from utils import factory
from trainer import _set_device, _set_random
import torch
import cv2  # OpenCV for image processing
import numpy as np  # NumPy for numerical operations
import timm  # PyTorch image models library


arg_list = [
    "--config",
    "exps/sdm/imr_inc10.json",
    "--prefix",
    "debug",
    "--device",
    "0",
]
arg_list = [
    "--config",
    "exps/baseline/single/cifar_inc10.json",
    "--prefix",
    "debug",
    "--device",
    "9",
]
a = setup_parser().parse_args(arg_list)
param = load_json(a.config)
args = merge_config(a, param)


args["train_seed"] = args["seed"][0]
_set_random(args["train_seed"])
_set_device(args)
data_manager = DataManager(
    args["dataset"],
    args["shuffle"],
    args["seed"],
    args["init_cls"],
    args["inc_cls"],
    args,
)

args["nb_classes"] = data_manager.nb_classes  # update args
args["nb_tasks"] = data_manager.nb_tasks
learner = get_learner(args["learner_name"], args)

tgt_task = 0
inc_cls = args["inc_cls"]
_known_classes = tgt_task * inc_cls
_total_classes = (tgt_task + 1) * inc_cls
test_dataset = data_manager.get_dataset(
    np.arange(0, _total_classes), source="test", mode="test"
)
test_loader = DataLoader(
    test_dataset,  # type: ignore
    batch_size=args["batch_size"],
    shuffle=False,
    num_workers=4,
)
_, x, y = next(iter(test_loader))
# display image x[idx]

idx = 3
input_img = x[idx].unsqueeze(0).to(learner._device)
target = y[idx].unsqueeze(0).to(learner._device)
model = torch.load(
    # "ckps/baseline/vit_base_patch16_224_in21k_baseSingle"
    # + "/imagenetr/20241011161227/imagenetr_0.pth"
    # "ckps/sdm/vit_base_patch16_224_in21k_sdm"
    # + "/imagenetr/20241011162216/imagenetr_0.pth"
    "ckps/baseline/vit_base_patch16_224_in21k_baseSingle/cifar224/20241024160026/cifar224_0.pth"
)
model.to(learner._device)
model.eval()


class GradCam:
    def __init__(self, model, tgt_layer):
        self.model = model.eval()  # Set the model to evaluation mode
        self.feature = None  # To store the features from the target layer
        self.gradient = None  # To store the gradients from the target layer
        self.handlers = []  # List to keep track of hooks
        self.tgt_layer = tgt_layer  # Target layer for Grad-CAM
        self._get_hook()  # Register hooks to the target layer

    # Hook to get features from the forward pass
    def _get_features_hook(self, module, input, output):
        self.feature = self.reshape_transform(
            output
        )  # Store and reshape the output features

    # Hook to get gradients from the backward pass
    def _get_grads_hook(self, module, input_grad, output_grad):
        self.gradient = self.reshape_transform(
            output_grad
        )  # Store and reshape the output gradients

        def _store_grad(grad):
            self.gradient = self.reshape_transform(
                grad
            )  # Store gradients for later use

        output_grad.register_hook(
            _store_grad
        )  # Register hook to store gradients

    # Register forward hooks to the target layer
    def _get_hook(self):
        self.tgt_layer.register_forward_hook(self._get_features_hook)
        self.tgt_layer.register_forward_hook(self._get_grads_hook)

    # Function to reshape the tensor for visualization
    def reshape_transform(self, tensor, height=14, width=14):
        result = tensor[:, 1:, :].reshape(
            tensor.size(0), height, width, tensor.size(2)
        )
        result = result.transpose(2, 3).transpose(
            1, 2
        )  # Rearrange dimensions to (C, H, W)
        return result

    # Function to compute the Grad-CAM heatmap
    def __call__(self, inputs):
        self.model.zero_grad()  # Zero the gradients
        # output = self.model(inputs)['logits']  # Forward pass
        output = self.model(inputs)  # Forward pass

        # Get the index of the highest score in the output
        index = np.argmax(output.cpu().data.numpy())
        target = output[0][index]  # Get the target score
        target.backward()  # Backward pass to compute gradients

        # Get the gradients and features
        gradient = self.gradient[0].cpu().data.numpy()  # type: ignore
        weight = np.mean(gradient, axis=(1, 2))  # Average the gradients
        feature = self.feature[0].cpu().data.numpy()  # type: ignore

        # Compute the weighted sum of the features
        cam = feature * weight[:, np.newaxis, np.newaxis]
        cam = np.sum(cam, axis=0)  # Sum over the channels
        cam = np.maximum(cam, 0)  # Apply ReLU to remove negative values

        # Normalize the heatmap
        cam -= np.min(cam)
        cam /= np.max(cam)
        cam = cv2.resize(
            cam, (224, 224)
        )  # Resize to match the input image size
        return cam  # Return the Grad-CAM heatmap


# Function to generate a Grad-CAM heatmap
def gen_cam(image, mask):
    # Create a heatmap from the Grad-CAM mask
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)  # type: ignore
    heatmap = np.float32(heatmap) / 255

    # Superimpose the heatmap on the original image
    cam = (1 - 0.5) * heatmap + 0.5 * (image[0]).cpu().permute(1, 2, 0).numpy()
    cam = cam / np.max(cam)  # Normalize the result
    return np.array(255 * cam, dtype=np.uint8)  # Convert to 8-bit image


target_layer = model.backbone.blocks[
    -1
].norm1  # Specify the target layer for Grad-CAM
model = timm.create_model("vit_base_patch16_224.augreg_in21k", pretrained=True)
target_layer = model.blocks[-1].norm1  # Specify the target layer for Grad-CAM
model.to(learner._device)
for param in model.parameters():
    param.requires_grad = True  # Disable gradient computation for the

# Initialize Grad-CAM with the model and target layer
grad_cam = GradCam(model, target_layer)
mask = grad_cam(input_img)  # Compute the Grad-CAM mask
result = gen_cam(input_img, mask)  # Generate the Grad-CAM heatmap

# Save the result to an image file
cv2.imwrite("result.jpg", result)
cv2.imwrite("img.jpg", (255 * input_img[0]).cpu().permute(1, 2, 0).numpy())
